package network

/*
网络流量监控系统 - 完整版
功能特性：
1. 多网卡自动检测
2. 原子操作保证并发安全
3. 滑动窗口速率平滑
4. 非活跃连接自动清理
注意，本方法执行使用了libpcap-dev库，需要先安装：
Linux (Debian/Ubuntu)：sudo apt-get install libpcap-dev
Linux (Fedora/RHEL/CentOS)：sudo yum install -y libpcap-dev
Linux (Debian/Ubuntu)：sudo apt-get install libpcap-dev
*/

import (
	"errors"
	"fmt"
	"net"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

const (
	windowSize     = 5
	cleanupTimeout = 5 * time.Minute
	refreshRate    = 200 // 毫秒[3](@ref)
)

type TrafficStat struct {
	totalBytes  uint64             // 累计总字节数[1](@ref)
	lastBytes   uint64             // 上次统计字节数[1](@ref)
	rateWindow  [windowSize]uint64 // 滑动窗口[3](@ref)
	windowIndex int32              // 窗口索引[3](@ref)
	lastUpdated time.Time          // 最后更新时间[4](@ref)
}

var (
	stats    sync.Map
	serverIP string
	iface    string
)

func getNetworkInfo() (string, string, error) {
	interfaces, err := net.Interfaces()
	if err != nil {
		return "", "", err
	}

	for _, i := range interfaces {
		if i.Flags&net.FlagUp == 0 || i.Flags&net.FlagLoopback != 0 {
			continue
		}

		addrs, _ := i.Addrs()
		for _, addr := range addrs {
			if ipnet, ok := addr.(*net.IPNet); ok && ipnet.IP.To4() != nil {
				return ipnet.IP.String(), i.Name, nil
			}
		}
	}
	return "", "", fmt.Errorf("no valid interface found")
}

func updateStats(clientIP string, packetLen uint) {

	val, _ := stats.LoadOrStore(clientIP, &TrafficStat{})
	stat := val.(*TrafficStat)

	atomic.AddUint64(&stat.totalBytes, uint64(packetLen))

	now := time.Now()
	elapsed := time.Since(stat.lastUpdated).Seconds()
	if elapsed >= float64(refreshRate)/1000 {
		current := atomic.LoadUint64(&stat.totalBytes)
		prev := atomic.SwapUint64(&stat.lastBytes, current)
		rate := uint64(float64(current-prev) / elapsed / 1024) // 速率计算[3](@ref)

		idx := atomic.AddInt32(&stat.windowIndex, 1) % windowSize
		stat.rateWindow[idx] = rate
		stat.lastUpdated = now
	}
}

func processPacket(packet gopacket.Packet) {
	networkLayer := packet.NetworkLayer()
	if networkLayer == nil {
		return
	}

	var srcIP, dstIP net.IP
	switch layer := networkLayer.(type) { // 协议解析[6](@ref)
	case *layers.IPv4:
		srcIP, dstIP = layer.SrcIP, layer.DstIP
	case *layers.IPv6:
		srcIP, dstIP = layer.SrcIP, layer.DstIP
	default:
		return
	}

	clientIP := ""
	switch {
	case srcIP.String() == serverIP:
		clientIP = dstIP.String()
	case dstIP.String() == serverIP:
		clientIP = srcIP.String()
	default:
		return
	}

	updateStats(clientIP, uint(packet.Metadata().Length))
}

func startCleanup() {
	go func() { // 自动清理协程[4](@ref)
		for range time.Tick(cleanupTimeout) {
			stats.Range(func(key, value interface{}) bool {
				stat := value.(*TrafficStat)
				if time.Since(stat.lastUpdated).Seconds() > cleanupTimeout.Seconds() {
					stats.Delete(key)
				}
				return true
			})
		}
	}()
}

// GetRate 获取指定IP的平均速率和总流量,AvgRate 平均速率(单位为KB/s),Total 总量(单位为MB)
func GetRate(ip string) (AvgRate float64, Total float64) {
	val, ok := stats.Load(ip)
	if !ok {
		return 0.0, 0.0
	}
	stat := val.(*TrafficStat)
	sum := uint64(0)
	for _, r := range stat.rateWindow {
		sum += r
	}
	avgRate := sum / windowSize

	return float64(avgRate), float64(atomic.LoadUint64(&stat.totalBytes) / 1024 / 1024)
}

// GetRateMb 获取指定IP的平均速率和总流量,AvgRate 平均速率(单位为MB/s),Total 总量(单位为GB)
func GetRateMb(ip string) (AvgRate float64, Total float64) {
	rate, total := GetRate(ip)
	return rate / 1024, total / 1024
}

func StartGetNetworkRate() (*pcap.Handle, error) {
	if os.Geteuid() != 0 { // 权限检查[1](@ref)
		return nil, errors.New("require sudo privileges")
	}

	var err error
	serverIP, iface, err = getNetworkInfo()
	if err != nil {
		return nil, err
	}

	handle, err := pcap.OpenLive(iface, 1600, true, pcap.BlockForever) // 抓包配置[6](@ref)
	if err != nil {
		return nil, err
	}

	if err := handle.SetBPFFilter(fmt.Sprintf("host %s", serverIP)); err != nil { // 过滤器设置[4](@ref)
		handle.Close()
		return nil, err
	}

	startCleanup()

	packetSource := gopacket.NewPacketSource( // 正确参数顺序[6](@ref)
		handle,
		handle.LinkType(),
	).Packets()

	go func() { // 异步处理协程[1](@ref)
		for packet := range packetSource {
			processPacket(packet)
		}
	}()
	return handle, nil
}
